#!/bin/env python
import os
import time

import cv2
import numpy as np
import gzip

import digits_ann as ANN

def inside(r1, r2):
    x1, y1, w1, h1 = r1
    x2, y2, w2, h2 = r2
    if (x1 > x2) and (y1 > y2) and (x1 + w1 < x2 + w2) and (y1 + h1 < y2 + h2):
        return True
    else:
        return False

def wrap_digit(rect):
    x, y, w, h = rect
    padding = 5
    hcenter = x + w / 2
    vcenter = y + h / 2
    if (h > w):
        w = h
        x = hcenter - (w / 2)
    else:
        h = w
        y = vcenter - (h / 2)

    return (x - padding, y - padding, w + padding, h + padding)

model_path = 'digital_recognition_model.xml'
if not os.path.exists(model_path):
    print(f"模型 {model_path} 不存在。")
    exit(-1)

# 加载训练好的ANN模型
ann = cv2.ml.ANN_MLP_load(model_path)

path = "../../resource/pic/digital.jpg"
img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)

# 高斯模糊，平滑图像
op1 = cv2.GaussianBlur(img, (3, 3), 0)

# 全局二值化, 反相图像颜色
ret, op2 = cv2.threshold(op1, 127, 255, cv2.THRESH_BINARY_INV)

# 腐蚀
op3 = cv2.erode(op2, np.ones((3, 3), np.uint8), iterations=1)

# 查找二值图轮廓
image, cntrs, hier = cv2.findContours(op3.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

rectangles = []

img2 = img.copy()

for c in cntrs:
    # 查找轮廓矩形
    r = x,y,w,h = cv2.boundingRect(c)
    # 获取轮廓面积
    a = cv2.contourArea(c)
    b = (img.shape[0] -3) * (img.shape[1] -3)

    is_inside = False
    for q in rectangles:
        if inside(r, q):
            is_inside = True
            break
    if not is_inside:
        if not a == b:
            rectangles.append(r)

    for r in rectangles:
        x, y, w, h = wrap_digit(r)
        # 将浮点坐标转换为整数
        x, y = int(x), int(y)  # 确保x和y是整数
        w, h = int(w), int(h)  # 确保w和h是整数
        cv2.rectangle(img2, (x, y), (x + w, y + h), (0, 255, 0), 2)
        roi = op3[y:y+h, x:x+w]

        try:
            digit_class = int(ANN.predict(ann, roi.copy())[0])
        except:
            continue
        cv2.putText(img2, "%d" % digit_class, (x, y-1), cv2.FONT_HERSHEY_SCRIPT_SIMPLEX, 1, (0, 255, 0))

show_img = np.hstack((img, op1, img2))
windowName = "digits ANN"
cv2.namedWindow(windowName, cv2.WINDOW_NORMAL)
cv2.resizeWindow(windowName, show_img.shape[1] * 2, show_img.shape[0] * 2)
cv2.imshow(windowName, show_img)

# show_img2 = np.hstack((op2, op3))
# windowName2 = "threshold"
# cv2.namedWindow(windowName2, cv2.WINDOW_NORMAL)
# cv2.resizeWindow(windowName2, show_img2.shape[1] * 2, show_img2.shape[0] * 2)
# cv2.imshow(windowName2, show_img2)

if __name__ == '__main__':
    while True:
      if cv2.waitKey(0) == ord('q'):
          break

    print('销毁窗口')
    cv2.destroyAllWindows()
